import pygame
import enum


class LinearGradientFillDirection(enum.Enum):
    HORIZONTAL = 0
    VERTICAL = 1
    MAIN_DIAGONAL = 2
    SECONDARY_DIAGONAL = 3


def linear_color(start: pygame.Color, end: pygame.Color, step: int, length: int) -> pygame.Color:
    def linear(start: int, end: int, step: int, length: int) -> int:
        if start == end or step < 0: return start
        if step > length: return end
        return int(step / length * (end - start) + start)

    return pygame.Color(
        linear(start.r, end.r, step, length),
        linear(start.g, end.g, step, length),
        linear(start.b, end.b, step, length),
        linear(start.a, end.a, step, length)
    )


def fill_linear_gradient_color(surface: pygame.Surface, start_color: pygame.Color, end_color: pygame.Color, direction: LinearGradientFillDirection) -> None:
    if start_color == end_color: return
    w, h = surface.get_size()

    if direction == LinearGradientFillDirection.HORIZONTAL:
        for i in range(w + 1):
            color = linear_color(start_color, end_color, i, w)
            pygame.draw.line(surface, color, (i, 0), (i, h))

    elif direction == LinearGradientFillDirection.VERTICAL:
        for i in range(h + 1):
            color = linear_color(start_color, end_color, i, h)
            pygame.draw.line(surface, color, (0, i), (w, i))

    elif direction == LinearGradientFillDirection.MAIN_DIAGONAL:
        length = max(w, h)
        for i in range(length + 1):
            color = linear_color(start_color, end_color, i, 2 * length)
            pygame.draw.line(surface, color, (0, i), ((i * w / length), 0))
        
        for i in range(length, 2 * length + 1):
            color = linear_color(start_color, end_color, i, 2 * length)
            pygame.draw.line(surface, color, (w, i - length), ((i - length) * w / length, length))

    elif direction == LinearGradientFillDirection.SECONDARY_DIAGONAL:
        length = max(w, h)
        for i in range(length + 1):
            color = linear_color(start_color, end_color, i, 2 * length)
            pygame.draw.line(surface, color, (0, length - i), ((i * w / length), length))
        
        for i in range(length, 2 * length + 1):
            color = linear_color(start_color, end_color, i, 2 * length)
            pygame.draw.line(surface, color, (w, 2 * length - i), ((length - i) * w / -length, 0))


def fill_radial_gradient_color(surface: pygame.Surface, start_color: pygame.Color, end_color: pygame.Color, radius: int, center: pygame.Vector2 = pygame.Vector2(0.5, 0.5)) -> None:
    if start_color == end_color or radius == 0: return

    center.x = center.x if center.x <= 1.0 else 1.0
    center.y = center.y if center.y <= 1.0 else 1.0
    center.x = 0.0 if center.x <= 0.0 else center.x
    center.y = 0.0 if center.y <= 0.0 else center.y

    w, h = surface.get_size()
    c_pos = pygame.Vector2()
    c_pos.x = center.x * w
    c_pos.y = center.y * h

    get_ellipse_rect = lambda x: (0, 0, w, h)
    if w >= h:
        scale = w / h
        get_ellipse_rect = lambda x: (c_pos.x - x, c_pos.y - x / scale, 2 * x, 2 * x / scale)
    else:
        scale = h / w
        get_ellipse_rect = lambda x: (c_pos.x - x / scale, c_pos.y - x, 2 * x / scale, 2 * x)

    
    surface.fill(end_color)
    for i in range(radius, 0, -1):
        color = linear_color(start_color, end_color, i, radius)
        pygame.draw.ellipse(surface, color, get_ellipse_rect(i), 2)

